/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.rocketmq.client.impl.consumer;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import org.apache.commons.lang3.StringUtils;
import org.apache.rocketmq.client.consumer.DefaultMQPushConsumer;
import org.apache.rocketmq.common.message.MessageAccessor;
import org.apache.rocketmq.common.message.MessageConst;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.remoting.protocol.body.ProcessQueueInfo;
import org.apache.rocketmq.logging.org.slf4j.Logger;
import org.apache.rocketmq.logging.org.slf4j.LoggerFactory;

/**
 * Queue consumption snapshot
 */
public class ProcessQueue {

    /**
     * 在执行再平衡的时候，处理队列加锁的超时时间：默认30秒
     */
    public final static long REBALANCE_LOCK_MAX_LIVE_TIME =
        Long.parseLong(System.getProperty("rocketmq.client.rebalance.lockMaxLiveTime", "30000"));

    /**
     * 在执行再平衡的时候，处理队列加锁的时间间隔：默认20秒
     */
    public final static long REBALANCE_LOCK_INTERVAL = Long.parseLong(System.getProperty("rocketmq.client.rebalance.lockInterval", "20000"));

    /**
     * 拉取消息操作的空闲时间：默认2分钟
     */
    private final static long PULL_MAX_IDLE_TIME = Long.parseLong(System.getProperty("rocketmq.client.pull.pullMaxIdleTime", "120000"));
    private final Logger log = LoggerFactory.getLogger(ProcessQueue.class);

    /**
     * 读写锁，用于对msgTreeMap执行加锁释放锁的操作，防止线程并发
     */
    private final ReadWriteLock treeMapLock = new ReentrantReadWriteLock();

    /**
     * 在消费者本地缓存的临时消息msgTreeMap
     * key：消息在MQ中的偏移量
     * value：消息实例
     */
    private final TreeMap<Long, MessageExt> msgTreeMap = new TreeMap<>();

    /**
     * 当前临时消息TreeMap中消息的总个数
     */
    private final AtomicLong msgCount = new AtomicLong();

    /**
     * 当前临时消息TreeMap中消息体的字节总长度
     */
    private final AtomicLong msgSize = new AtomicLong();

    /**
     * 消费加锁，重入互斥锁
     */
    private final Lock consumeLock = new ReentrantLock();

    /**
     * msgTreeMap的子集，仅在有序消费的时候用到
     */
    private final TreeMap<Long, MessageExt> consumingMsgOrderlyTreeMap = new TreeMap<>();

    /**
     * 尝试解锁次数
     */
    private final AtomicLong tryUnlockTimes = new AtomicLong(0);

    /**
     * 当前处理队列中消息的MQ偏移量的最大值
     */
    private volatile long queueOffsetMax = 0L;
    private volatile boolean dropped = false;
    private volatile long lastPullTimestamp = System.currentTimeMillis();
    private volatile long lastConsumeTimestamp = System.currentTimeMillis();

    /**
     * 处理队列是否被锁定
     */
    private volatile boolean locked = false;
    private volatile long lastLockTimestamp = System.currentTimeMillis();

    /**
     * 当前处理队列是否在消费中
     */
    private volatile boolean consuming = false;

    /**
     * 当前处理队列中积压的消息总数
     */
    private volatile long msgAccCnt = 0;

    /**
     * 是否加锁超时
     * @return
     */
    public boolean isLockExpired() {
        // 从最后一次加锁到现在的时间 是否 大于 再平衡加锁最长时间：默认30秒
        return (System.currentTimeMillis() - this.lastLockTimestamp) > REBALANCE_LOCK_MAX_LIVE_TIME;
    }

    /**
     * 是否拉取超时
     * @return
     */
    public boolean isPullExpired() {
        // 从最后一次拉取到现在的时间 是否 大于 拉取闲置时间：默认两分钟
        return (System.currentTimeMillis() - this.lastPullTimestamp) > PULL_MAX_IDLE_TIME;
    }

    /**
     * 清理过期消息
     * @param pushConsumer 对应的消费者实例
     */
    public void cleanExpiredMsg(DefaultMQPushConsumer pushConsumer) {
        // 如果消费者是有序消费，则直接返回，不处理
        if (pushConsumer.isConsumeOrderly()) {
            return;
        }
        // 最多处理16条消息，多了也没用，因为只需要查看第一条消息是否处理超时，如果超时，移除，再次查看第一条消息是否处理超时
        int loop = Math.min(msgTreeMap.size(), 16);
        for (int i = 0; i < loop; i++) {
            MessageExt msg = null;
            try {
                // 获取读锁
                this.treeMapLock.readLock().lockInterruptibly();
                try {
                    if (!msgTreeMap.isEmpty()) {
                        // 获取开始消费的时间戳
                        String consumeStartTimeStamp = MessageAccessor.getConsumeStartTimeStamp(msgTreeMap.firstEntry().getValue());
                        // 如果时间戳存在，并且开始消费的时间到当前时间的时间差大于消费者消费超时时间
                        // 则消息消费过期
                        if (StringUtils.isNotEmpty(consumeStartTimeStamp) && System.currentTimeMillis() - Long.parseLong(consumeStartTimeStamp) > pushConsumer.getConsumeTimeout() * 60 * 1000) {
                            // 获取该消息
                            msg = msgTreeMap.firstEntry().getValue();
                        }
                    }
                } finally {
                    // 释放读锁
                    this.treeMapLock.readLock().unlock();
                }
            } catch (InterruptedException e) {
                log.error("getExpiredMsg exception", e);
            }

            // 如果没有找到消息，则跳出循环：最早的消息都没有超时，其他的就不用看了
            if (msg == null) {
                break;
            }

            try {
                // 消费者将消费超时的消息发回给broker
                pushConsumer.sendMessageBack(msg, 3);
                log.info("send expire msg back. topic={}, msgId={}, storeHost={}, queueId={}, queueOffset={}", msg.getTopic(), msg.getMsgId(), msg.getStoreHost(), msg.getQueueId(), msg.getQueueOffset());
                try {
                    // 获取写锁
                    this.treeMapLock.writeLock().lockInterruptibly();
                    try {
                        if (!msgTreeMap.isEmpty() && msg.getQueueOffset() == msgTreeMap.firstKey()) {
                            try {
                                // 将指定的消息从msgTreeMap中移除
                                removeMessage(Collections.singletonList(msg));
                            } catch (Exception e) {
                                log.error("send expired msg exception", e);
                            }
                        }
                    } finally {
                        // 释放写锁
                        this.treeMapLock.writeLock().unlock();
                    }
                } catch (InterruptedException e) {
                    log.error("getExpiredMsg exception", e);
                }
            } catch (Exception e) {
                log.error("send expired msg exception", e);
            }
        }
    }

    /**
     * 向处理队列中添加消息
     * @param msgs 需要添加到处理队列中的消息集合
     * @return 是否需要分配消费：如果处理队列有消息，且没有正在消费，则返回true。
     */
    public boolean putMessage(final List<MessageExt> msgs) {
        boolean dispatchToConsume = false;
        try {
            // 获取写锁
            this.treeMapLock.writeLock().lockInterruptibly();
            try {
                int validMsgCnt = 0;
                for (MessageExt msg : msgs) {
                    // 将消息的MQ偏移量为key，消息为value的键值对添加到消息map中
                    MessageExt old = msgTreeMap.put(msg.getQueueOffset(), msg);
                    // 如果返回值为null，表示是新的消息
                    if (null == old) {
                        // 累加添加到处理队列中的消息个数
                        validMsgCnt++;
                        // 处理队列中最大的MQ偏移量
                        this.queueOffsetMax = msg.getQueueOffset();
                        // 当前处理队列中消息的总字节长度累加当前新消息体的字节长度
                        msgSize.addAndGet(msg.getBody().length);
                    }
                }
                // 当前处理队列中消息的条目数累加
                msgCount.addAndGet(validMsgCnt);

                // 如果当前处理队列中有消息，但是并没有在处理，则设置标记：需要分发给消费处理
                if (!msgTreeMap.isEmpty() && !this.consuming) {
                    // 需要分发给消费处理
                    dispatchToConsume = true;
                    this.consuming = true;
                }

                // 如果参数不是空的
                if (!msgs.isEmpty()) {
                    // 获取最后一条消息
                    MessageExt messageExt = msgs.get(msgs.size() - 1);
                    // 获取 "MAX_OFFSET"
                    String property = messageExt.getProperty(MessageConst.PROPERTY_MAX_OFFSET);
                    if (property != null) {
                        // 计算当前处理队列积压的消息数
                        // 如果消息包含属性MAX_OFFSET属性，则使用该属性值减去当前最新消息的偏移量，得出在broker端积压的消息数
                        long accTotal = Long.parseLong(property) - messageExt.getQueueOffset();
                        if (accTotal > 0) {
                            this.msgAccCnt = accTotal;
                        }
                    }
                }
            } finally {
                // 释放写锁
                this.treeMapLock.writeLock().unlock();
            }
        } catch (InterruptedException e) {
            log.error("putMessage exception", e);
        }

        // 返回标记：是否需要执行消费处理
        return dispatchToConsume;
    }

    /**
     * 当前处理队列中消息的MQ偏移量跨度
     * 当前处理队列中最后一条消息的MQ偏移量减去处理队列中第一条消息的MQ偏移量得到MQ偏移量跨度
     * @return
     */
    public long getMaxSpan() {
        try {
            // 获取读锁
            this.treeMapLock.readLock().lockInterruptibly();
            try {
                if (!this.msgTreeMap.isEmpty()) {
                    // 获取最后一个消息的MQ偏移量 减去 第一个消息的MQ偏移量，得到消息MQ偏移量跨度
                    return this.msgTreeMap.lastKey() - this.msgTreeMap.firstKey();
                }
            } finally {
                // 释放读锁
                this.treeMapLock.readLock().unlock();
            }
        } catch (InterruptedException e) {
            log.error("getMaxSpan exception", e);
        }

        return 0;
    }

    /**
     * 当消费者将消费超时的消息发回给broker后，删除本地缓存的消息
     * @param msgs 需要删除的消费超时的消息集合
     * @return
     */
    public long removeMessage(final List<MessageExt> msgs) {
        long result = -1;
        final long now = System.currentTimeMillis();
        try {
            // 获取写锁
            this.treeMapLock.writeLock().lockInterruptibly();
            this.lastConsumeTimestamp = now;
            try {
                if (!msgTreeMap.isEmpty()) {
                    result = this.queueOffsetMax + 1;
                    int removedCnt = 0;
                    for (MessageExt msg : msgs) {
                        // 从treeMap中根据消息的MQ偏移量删除消息
                        MessageExt prev = msgTreeMap.remove(msg.getQueueOffset());
                        if (prev != null) {
                            // 累减
                            removedCnt--;
                            // 消息字节长度减去被删除消息的消息体长度
                            msgSize.addAndGet(0 - msg.getBody().length);
                        }
                    }
                    // 消息总数减去对应的值
                    msgCount.addAndGet(removedCnt);

                    if (!msgTreeMap.isEmpty()) {
                        // 赋值新的最早消息在MQ中的偏移量
                        result = msgTreeMap.firstKey();
                    }
                }
            } finally {
                this.treeMapLock.writeLock().unlock();
            }
        } catch (Throwable t) {
            log.error("removeMessage exception", t);
        }

        // 返回处理队列中新的最早的消息
        return result;
    }

    public TreeMap<Long, MessageExt> getMsgTreeMap() {
        return msgTreeMap;
    }

    public AtomicLong getMsgCount() {
        return msgCount;
    }

    public AtomicLong getMsgSize() {
        return msgSize;
    }

    public boolean isDropped() {
        return dropped;
    }

    public void setDropped(boolean dropped) {
        this.dropped = dropped;
    }

    public boolean isLocked() {
        return locked;
    }

    public void setLocked(boolean locked) {
        this.locked = locked;
    }

    /**
     * 回滚
     */
    public void rollback() {
        try {
            // 获取写锁
            this.treeMapLock.writeLock().lockInterruptibly();
            try {
                // 将正在处理中的有序消费对应的Map中的消息添加到处理队列中
                this.msgTreeMap.putAll(this.consumingMsgOrderlyTreeMap);
                // 将正在处理的有序消息的TreeMap清空
                this.consumingMsgOrderlyTreeMap.clear();
            } finally {
                // 释放写锁
                this.treeMapLock.writeLock().unlock();
            }
        } catch (InterruptedException e) {
            log.error("rollback exception", e);
        }
    }

    /**
     * 提交
     * @return
     */
    public long commit() {
        try {
            // 获取写锁
            this.treeMapLock.writeLock().lockInterruptibly();
            try {
                // 获取正在执行有序消费的TreeMap中最后一个消息的MQ偏移量
                Long offset = this.consumingMsgOrderlyTreeMap.lastKey();
                // 消息总数减去正在处理的TreeMap中的消息个数
                msgCount.addAndGet(0 - this.consumingMsgOrderlyTreeMap.size());
                // 遍历正在消费的TreeMap
                for (MessageExt msg : this.consumingMsgOrderlyTreeMap.values()) {
                    // 消息字节总大小挨个减去其中消息体的长度
                    msgSize.addAndGet(0 - msg.getBody().length);
                }
                // 将正在有序处理的TreeMap中的消息清空
                this.consumingMsgOrderlyTreeMap.clear();

                // 如果有消息，则偏移量+1，得出下一条该处理的消息偏移量
                if (offset != null) {
                    return offset + 1;
                }
            } finally {
                // 释放写锁
                this.treeMapLock.writeLock().unlock();
            }
        } catch (InterruptedException e) {
            log.error("commit exception", e);
        }
        return -1;
    }

    /**
     * 将消息从正在有序处理的队列移到消息的待处理TreeMap中
     * @param msgs 需要重新进入待处理TreeMap中的消息集合
     */
    public void makeMessageToConsumeAgain(List<MessageExt> msgs) {
        try {
            // 获取写锁
            this.treeMapLock.writeLock().lockInterruptibly();
            try {
                for (MessageExt msg : msgs) {
                    // 移除
                    this.consumingMsgOrderlyTreeMap.remove(msg.getQueueOffset());
                    // 加入
                    this.msgTreeMap.put(msg.getQueueOffset(), msg);
                }
            } finally {
                // 释放写锁
                this.treeMapLock.writeLock().unlock();
            }
        } catch (InterruptedException e) {
            log.error("makeMessageToCosumeAgain exception", e);
        }
    }

    /**
     * 将msgTreeMap中的消息移动到结果集合中，同时也添加到正在有序处理的TreeMap中
     * @param batchSize 需要最多移动的消息条数
     * @return 从msgTreeMap中移出来的消息集合
     */
    public List<MessageExt> takeMessages(final int batchSize) {
        // 结果容器
        List<MessageExt> result = new ArrayList<>(batchSize);
        final long now = System.currentTimeMillis();
        try {
            // 获取写锁
            this.treeMapLock.writeLock().lockInterruptibly();
            this.lastConsumeTimestamp = now;
            try {
                if (!this.msgTreeMap.isEmpty()) {
                    // 最多获取batchSize条消息
                    for (int i = 0; i < batchSize; i++) {
                        // 拉取第一条消息
                        Map.Entry<Long, MessageExt> entry = this.msgTreeMap.pollFirstEntry();
                        if (entry != null) {
                            // 添加到结果集合中
                            result.add(entry.getValue());
                            // 添加到正在有序处理的TreeMap中
                            consumingMsgOrderlyTreeMap.put(entry.getKey(), entry.getValue());
                        } else {
                            break;
                        }
                    }
                }

                // 如果结果集合是空的，设置标志：没有正在消费
                if (result.isEmpty()) {
                    consuming = false;
                }
            } finally {
                // 释放写锁
                this.treeMapLock.writeLock().unlock();
            }
        } catch (InterruptedException e) {
            log.error("take Messages exception", e);
        }

        // 返回结果
        return result;
    }

    /**
     * Return the result that whether current message is exist in the process queue or not.
     * 查询指定的消息是否在msgTreeMap中。
     * @param message 要查询的消息
     * @return true表示包含，false表示不包含
     */
    public boolean containsMessage(MessageExt message) {
        if (message == null) {
            // should never reach here.
            return false;
        }
        try {
            // 获取读锁
            this.treeMapLock.readLock().lockInterruptibly();
            try {
                // msgTreeMap中是否包含指定的消息
                return this.msgTreeMap.containsKey(message.getQueueOffset());
            } finally {
                // 释放读锁
                this.treeMapLock.readLock().unlock();
            }
        } catch (Throwable t) {
            log.error("Failed to check message's existence in process queue, message={}", message, t);
        }
        return false;
    }

    /**
     * 是否有临时消息：msgTreeMap是否为空的，空的返回false，不空返回true。
     * @return true表示有，false表示没有
     */
    public boolean hasTempMessage() {
        try {
            // 获取读锁
            this.treeMapLock.readLock().lockInterruptibly();
            try {
                // msgTreeMap是否为空的
                return !this.msgTreeMap.isEmpty();
            } finally {
                // 释放读锁
                this.treeMapLock.readLock().unlock();
            }
        } catch (InterruptedException e) {
        }

        return true;
    }

    /**
     * 清空msgTreeMap；
     * 清空正在有序消费的TreeMap；
     * 将消息个数设置为0；
     * 将消息体总长度设置为0；
     * 将MQ消息偏移量最大值设置为0。
     */
    public void clear() {
        try {
            // 获取写锁
            this.treeMapLock.writeLock().lockInterruptibly();
            try {
                // 清空msgTreeMap
                this.msgTreeMap.clear();
                // 清空正在有序消费的TreeMap
                this.consumingMsgOrderlyTreeMap.clear();
                // 将消息个数设置为0
                this.msgCount.set(0);
                // 将消息体总长度设置为0
                this.msgSize.set(0);
                // 将MQ消息偏移量最大值设置为0
                this.queueOffsetMax = 0L;
            } finally {
                // 释放写锁
                this.treeMapLock.writeLock().unlock();
            }
        } catch (InterruptedException e) {
            log.error("rollback exception", e);
        }
    }

    public long getLastLockTimestamp() {
        return lastLockTimestamp;
    }

    public void setLastLockTimestamp(long lastLockTimestamp) {
        this.lastLockTimestamp = lastLockTimestamp;
    }

    public Lock getConsumeLock() {
        return consumeLock;
    }

    public long getLastPullTimestamp() {
        return lastPullTimestamp;
    }

    public void setLastPullTimestamp(long lastPullTimestamp) {
        this.lastPullTimestamp = lastPullTimestamp;
    }

    public long getMsgAccCnt() {
        return msgAccCnt;
    }

    public void setMsgAccCnt(long msgAccCnt) {
        this.msgAccCnt = msgAccCnt;
    }

    public long getTryUnlockTimes() {
        return this.tryUnlockTimes.get();
    }

    public void incTryUnlockTimes() {
        this.tryUnlockTimes.incrementAndGet();
    }

    /**
     * 将当前处理队列的信息填充给参数
     * @param info
     */
    public void fillProcessQueueInfo(final ProcessQueueInfo info) {
        try {
            // 对msgTreeMap加读锁
            this.treeMapLock.readLock().lockInterruptibly();

            if (!this.msgTreeMap.isEmpty()) {
                // 设置本地缓存的消息的最小偏移量
                info.setCachedMsgMinOffset(this.msgTreeMap.firstKey());
                // 设置本地缓存的消息的最大偏移量
                info.setCachedMsgMaxOffset(this.msgTreeMap.lastKey());
                // 设置本地缓存的消息总数
                info.setCachedMsgCount(this.msgTreeMap.size());
                // 设置本地缓存的消息体字节总长度，单位MB
                info.setCachedMsgSizeInMiB((int) (this.msgSize.get() / (1024 * 1024)));
            }

            // 如果是有序消费，则
            if (!this.consumingMsgOrderlyTreeMap.isEmpty()) {
                // 设置事务消息最小偏移量
                info.setTransactionMsgMinOffset(this.consumingMsgOrderlyTreeMap.firstKey());
                // 设置事务消息最大偏移量
                info.setTransactionMsgMaxOffset(this.consumingMsgOrderlyTreeMap.lastKey());
                // 设置本地缓存的事务消息总数
                info.setTransactionMsgCount(this.consumingMsgOrderlyTreeMap.size());
            }

            // 设置是否被锁定
            info.setLocked(this.locked);
            // 设置尝试解锁次数
            info.setTryUnlockTimes(this.tryUnlockTimes.get());
            // 设置最后一次加锁时间戳
            info.setLastLockTimestamp(this.lastLockTimestamp);

            // 设置是否被丢弃
            info.setDroped(this.dropped);
            // 设置最后一次拉取消息时间戳
            info.setLastPullTimestamp(this.lastPullTimestamp);
            // 设置最后一次消费时间戳
            info.setLastConsumeTimestamp(this.lastConsumeTimestamp);
        } catch (Exception e) {
        } finally {
            // 释放读锁
            this.treeMapLock.readLock().unlock();
        }
    }

    public long getLastConsumeTimestamp() {
        return lastConsumeTimestamp;
    }

    public void setLastConsumeTimestamp(long lastConsumeTimestamp) {
        this.lastConsumeTimestamp = lastConsumeTimestamp;
    }

}
